package com.maven.user.util;

import org.dbunit.DatabaseUnitException;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.database.QueryDataSet;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSet;
import org.dbunit.dataset.xml.FlatXmlDataSetBuilder;
import org.dbunit.ext.mysql.MySqlConnection;
import org.dbunit.operation.DatabaseOperation;
import org.junit.AfterClass;
import org.junit.Assert;
import org.junit.BeforeClass;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.io.InputStream;
import java.sql.SQLException;

public abstract class AbstractDbUnitTestCase {
    public static IDatabaseConnection dbConnection;
    private File tempFile;

    @BeforeClass
    public static void init() throws DatabaseUnitException {
        //这里Mysql要加上schema，不然会报错，而且这里要改成MySqlConnection
        dbConnection = new MySqlConnection(DbUtil.getConnection(), "test");
    }

    protected IDataSet createDateSet(String tName) throws DataSetException {
        InputStream in = AbstractDbUnitTestCase.class.getClassLoader()
                .getResourceAsStream(tName + ".xml");
        Assert.assertNotNull("dbunit的基本数据文件不存在", in);
        return new FlatXmlDataSetBuilder().build(in);
    }


    protected void backupAllTable() throws SQLException, DataSetException, IOException {
        IDataSet ds = dbConnection.createDataSet();
        writeBackupFile(ds);
    }

    private void writeBackupFile(IDataSet ds) throws IOException, DataSetException {
        tempFile = new File("src/test/resources/backup.xml");
        FlatXmlDataSet.write(ds, new FileWriter(tempFile), "utf-8");
    }

    /**
     * 备份自定义表格
     *
     * @param tname
     * @throws DataSetException
     * @throws IOException
     */
    private void backupCustomTable(String[] tname) throws DataSetException, IOException {
        QueryDataSet ds = new QueryDataSet(dbConnection);
        for (String str : tname) {
            ds.addTable(str);
        }
        writeBackupFile(ds);
    }

    /**
     * 备份一个表格
     *
     * @param tname
     * @throws DataSetException
     * @throws IOException
     */
    protected void backupOneTable(String tname) throws DataSetException, IOException {
        backupCustomTable(new String[]{tname});
    }

    protected void resumeTable() throws DatabaseUnitException, SQLException, IOException {
        IDataSet ds = new FlatXmlDataSetBuilder().build(tempFile);
        DatabaseOperation.CLEAN_INSERT.execute(dbConnection, ds);
    }

    @AfterClass
    public static void destroy() {
        try {
            if (dbConnection != null) {
                dbConnection.close();
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
